import pandas as pd
import numpy as np
import statsmodels.api as sm
from sklearn.metrics import mean_absolute_error, mean_squared_error
from typing import Tuple, Optional, Dict, List
import warnings

class SARIMAXForecaster:
    """
    SARIMAX 季节性时间序列预测模型（支持多元外部变量）
    
    功能：
    1. 自动处理时间序列数据（缺失值填充、季节性差分）
    2. 支持外部回归变量（如经济指标、政策变量）
    3. 提供网格搜索优化参数（p,d,q,P,D,Q,s）
    4. 输出预测结果和模型评估指标
    """
    
    def __init__(self, order: Tuple[int, int, int] = (1, 0, 1),
                 seasonal_order: Tuple[int, int, int, int] = (1, 1, 1, 12),
                 trend: str = 'c',
                 enforce_stationarity: bool = True,
                 enforce_invertibility: bool = True):
        """
        初始化 SARIMAX 模型
        
        参数:
        - order: 非季节性 (p, d, q) 阶数
        - seasonal_order: 季节性 (P, D, Q, s) 阶数，s为季节周期（如月度数据s=12）
        - trend: 趋势项 ('c'=常数, 't'=线性, 'ct'=两者)
        - enforce_stationarity: 强制平稳性
        - enforce_invertibility: 强制可逆性
        """
        self.order = order
        self.seasonal_order = seasonal_order
        self.trend = trend
        self.enforce_stationarity = enforce_stationarity
        self.enforce_invertibility = enforce_invertibility
        self.model = None
        self.model_results = None
        self.best_params = None
        
    def fit(self, y: pd.Series, exog: Optional[pd.DataFrame] = None,
            suppress_warnings: bool = True) -> None:
        """
        拟合 SARIMAX 模型
        
        参数:
        - y: 目标时间序列（如就业率）
        - exog: 外部变量 DataFrame（需与y索引对齐）
        - suppress_warnings: 是否抑制statsmodels的收敛警告
        """
        if suppress_warnings:
            warnings.filterwarnings("ignore", category=UserWarning)
        
        # 检查数据
        if exog is not None:
            assert len(y) == len(exog), "y和exog长度必须一致"
            exog = exog.reindex(y.index)  # 对齐索引
        
        # 拟合模型
        self.model = sm.tsa.SARIMAX(
            endog=y,
            exog=exog,
            order=self.order,
            seasonal_order=self.seasonal_order,
            trend=self.trend,
            enforce_stationarity=self.enforce_stationarity,
            enforce_invertibility=self.enforce_invertibility
        )
        
        self.model_results = self.model.fit(disp=False)
        
    def predict(self, steps: int, exog_future: Optional[pd.DataFrame] = None) -> pd.Series:
        """
        预测未来steps步
        
        参数:
        - steps: 预测步长
        - exog_future: 未来steps步的外部变量（必须提供）
        
        返回:
        - 预测结果 Series（带时间索引）
        """
        if self.model_results is None:
            raise ValueError("模型未训练，请先调用fit()")
        
        if exog_future is not None:
            assert len(exog_future) == steps, "exog_future长度必须等于steps"
        
        # 生成预测
        forecast = self.model_results.get_forecast(
            steps=steps,
            exog=exog_future
        )
        
        return forecast.predicted_mean
    
    def grid_search(self, y: pd.Series, exog: Optional[pd.DataFrame] = None,
                    param_grid: Dict[str, List] = None,
                    cv_splits: int = 3) -> Dict:
        """
        网格搜索优化 SARIMAX 参数（基于时间序列交叉验证）
        
        参数:
        - y: 目标序列
        - exog: 外部变量
        - param_grid: 参数网格，例如：
            {
                'p': [0, 1, 2],
                'd': [0, 1],
                'q': [0, 1, 2],
                'P': [0, 1],
                'D': [0, 1],
                'Q': [0, 1],
                's': [6, 12]  # 季节周期
            }
        - cv_splits: 交叉验证折数
        
        返回:
        - 最佳参数和对应MAE
        """
        if param_grid is None:
            param_grid = {
                'p': [0, 1, 2],
                'd': [0, 1],
                'q': [0, 1, 2],
                'P': [0, 1],
                'D': [0, 1],
                'Q': [0, 1],
                's': [6, 12]  # 假设季节周期为6或12
            }
        
        from itertools import product
        best_mae = float('inf')
        best_params = None
        
        # 生成参数组合
        param_combinations = product(
            param_grid['p'], param_grid['d'], param_grid['q'],
            param_grid['P'], param_grid['D'], param_grid['Q'],
            param_grid['s']
        )
        
        # 时间序列交叉验证
        cv_size = len(y) // cv_splits
        for params in param_combinations:
            p, d, q, P, D, Q, s = params
            current_order = (p, d, q)
            current_seasonal_order = (P, D, Q, s)
            
            mae_scores = []
            for i in range(cv_splits):
                # 划分训练集和验证集
                val_start = i * cv_size
                val_end = (i + 1) * cv_size
                y_train = y.iloc[:val_start]
                y_val = y.iloc[val_start:val_end]
                
                exog_train = exog.iloc[:val_start] if exog is not None else None
                exog_val = exog.iloc[val_start:val_end] if exog is not None else None
                
                # 训练模型
                try:
                    model = sm.tsa.SARIMAX(
                        endog=y_train,
                        exog=exog_train,
                        order=current_order,
                        seasonal_order=current_seasonal_order,
                        enforce_stationarity=True,
                        enforce_invertibility=True
                    )
                    results = model.fit(disp=False)
                    
                    # 预测验证集
                    forecast = results.get_forecast(
                        steps=len(y_val),
                        exog=exog_val
                    )
                    pred = forecast.predicted_mean
                    mae = mean_absolute_error(y_val, pred)
                    mae_scores.append(mae)
                except:
                    continue
            
            # 更新最佳参数
            if mae_scores and np.mean(mae_scores) < best_mae:
                best_mae = np.mean(mae_scores)
                best_params = {
                    'order': current_order,
                    'seasonal_order': current_seasonal_order,
                    'mae': best_mae
                }
        
        self.best_params = best_params
        return best_params
    
    def evaluate(self, y_true: pd.Series, y_pred: pd.Series) -> Dict:
        """评估预测结果"""
        return {
            'MAE': mean_absolute_error(y_true, y_pred),
            'RMSE': np.sqrt(mean_squared_error(y_true, y_pred)),
            'MAPE': np.mean(np.abs((y_true - y_pred) / y_true)) * 100
        }
    
    def summary(self) -> str:
        """返回模型摘要"""
        if self.model_results is None:
            return "模型未训练"
        return self.model_results.summary().as_text()